import networkx as nx
import pulp
import matplotlib.pyplot as plt
import numpy as np
import random

# 1. 随机拓扑生成
def generate_network(num_nodes=30, num_leaks=5, num_sources=3):
    G = nx.Graph()
    
    # 添加节点
    for i in range(num_nodes):
        if i < num_sources:
            G.nodes[i]['type'] = 'source'
            G.nodes[i]['pressure'] = random.uniform(0.4, 0.6)  # MPa
        elif i < num_sources + num_leaks:
            G.nodes[i]['type'] = 'leak'
            G.nodes[i]['leak_rate'] = random.uniform(0.1, 1.0)  # m³/h
        else:
            G.nodes[i]['type'] = 'user'
            G.nodes[i]['importance'] = random.choice([1, 2, 3])  # 重要性分级
    
    # 添加边（管道）
    for _ in range(num_nodes*3//2):
        u, v = random.sample(range(num_nodes), 2)
        if not G.has_edge(u, v):
            length = random.uniform(50, 500)  # 米
            diameter = random.choice([50, 80, 100, 150, 200])  # mm
            roughness = 0.17  # mm (钢管)
            G.add_edge(u, v, length=length, diameter=diameter, 
                      roughness=roughness, valve=random.random()<0.1)

    return G

# 2. 最小割算法计算关阀节点
def min_cut_isolation(G, leaks, sources):
    cut_nodes = set()
    for leak in leaks:
        # 构建流网络
        flow_graph = nx.DiGraph()
        for u, v in G.edges:
            # 双向管道建模
            flow_graph.add_edge(u, v, capacity=1.0)
            flow_graph.add_edge(v, u, capacity=1.0)
        
        # 执行最小割
        cut_value, partition = nx.minimum_cut(flow_graph, leak, sources)
        reachable, non_reachable = partition
        cut_nodes.update(non_reachable)
    
    return list(cut_nodes)

# 3. 线性规划模型
def optimize_valve_closing(G, leaks, sources, max_time=30):
    prob = pulp.LpProblem("Valve_Closing_Optimization", pulp.LpMinimize)
    
    # 决策变量
    valve_vars = pulp.LpVariable.dicts("Valve", G.nodes, cat='Binary')
    time_vars = pulp.LpVariable.dicts("Time", G.nodes, lowBound=0)
    
    # 目标函数（加权求和）
    prob += 0.5*pulp.lpSum(valve_vars[v] for v in G.nodes if G.nodes[v].get('type')=='user') + \
            0.3*pulp.lpSum(G.edges[u,v]['length'] for u,v in G.edges if valve_vars[u] or valve_vars[v]) + \
            0.2*pulp.lpSum(time_vars[v] for v in G.nodes if valve_vars[v])
    
    # 约束条件
    for leak in leaks:
        prob += pulp.lpSum(valve_vars[v] for v in G.predecessors(leak)) >= 1
    
    for v in G.nodes:
        if G.nodes[v].get('type') == 'valve':
            prob += time_vars[v] <= max_time
    
    # 求解
    prob.solve(pulp.PULP_CBC_CMD(msg=0))
    
    return {v.name: pulp.value(var) for v, var in valve_vars.items()}

# 4. 可视化
def visualize_network(G, cut_nodes, valve_status):
    pos = nx.spring_layout(G)
    plt.figure(figsize=(12,8))
    
    # 绘制节点
    for node in G.nodes:
        if G.nodes[node]['type'] == 'source':
            nx.draw_networkx_nodes(G, pos, nodelist=[node], node_color='green', node_size=1500)
        elif G.nodes[node]['type'] == 'leak':
            nx.draw_networkx_nodes(G, pos, nodelist=[node], node_color='red', node_size=1000)
        else:
            nx.draw_networkx_nodes(G, pos, nodelist=[node], node_color='skyblue', node_size=800)
    
    # 绘制边
    edges = G.edges()
    lengths = [G.edges[u,v]['length'] for u,v in edges]
    diameters = [G.edges[u,v]['diameter'] for u,v in edges]
    
    # 边宽度映射
    edge_widths = [10 if G.nodes[u].get('type')=='valve' else 
                  (30 if v in cut_nodes else 15) for u,v in edges]
    
    nx.draw_networkx_edges(G, pos, edgelist=edges, width=edge_widths, alpha=0.5)
    
    # 标注关键信息
    labels = {node: f"{node}\n{G.nodes[node].get('type')}\n{G.nodes[node].get('pressure', '')}"
             for node in G.nodes}
    nx.draw_networkx_labels(G, pos, labels, font_size=8)
    
    # 标注关闭阀门
    close_nodes = [v for v in valve_status if valve_status[v]>0.5]
    nx.draw_networkx_nodes(G, pos, nodelist=close_nodes, node_color='black', node_size=1200)
    
    plt.title("燃气泄漏关阀优化方案")
    plt.axis('off')
    plt.show()

# 主程序
if __name__ == "__main__":
    # 参数设置
    np.random.seed(42)
    random.seed(42)
    
    # 生成网络
    G = generate_network(num_nodes=50, num_leaks=5, num_sources=3)
    
    # 获取泄漏点和气源
    leaks = [n for n in G.nodes if G.nodes[n]['type']=='leak']
    sources = [n for n in G.nodes if G.nodes[n]['type']=='source']
    
    # 最小割计算
    cut_nodes = min_cut_isolation(G, leaks, sources)
    
    # 线性规划优化
    valve_status = optimize_valve_closing(G, leaks, sources)
    
    # 可视化
    visualize_network(G, cut_nodes, valve_status)